import tensorflow as tf
import pandas as pd
import numpy as np
from Core.dataset import get_generator
from Core.Model import AlexNet
from tensorflow.keras.preprocessing.image import ImageDataGenerator

train_generator, _, test_generator = get_generator()

# 提取模型
model_save_path = './checkpoint/AlexNet.ckpt'
model = AlexNet()
model.load_weights(model_save_path)
# model.summary()

# predict_generator 以概率给出输出,因此首先我们需要将它们转换为类数,如0,1。
# 如果我们同时使用了两个生成器，我们就需要对生成器进行reset，否则文件名等会乱七八糟。
test_generator.reset()
pred = model.predict_generator(test_generator, verbose=1)

predicted_class_indices = np.argmax(pred, axis=1)
labels = train_generator.class_indices
label = dict((v, k) for k, v in labels.items())

# 建立代码标签与真实标签的关系
predictions = [label[i] for i in predicted_class_indices]

# 建立预测结果和文件名之间的关系
filenames = test_generator.filenames
# results = pd.DataFrame({"Filename": filenames,
#                         "Predictions": predictions})
# print(results)
for idx in range(len(filenames)):
    print('predict %s' % (predictions[idx]))
    print('title   %s' % filenames[idx])
    print('')
